package com.example.im.config;

import com.example.im.message.MessageManager;
import com.example.im.protocol.Message;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.net.URLDecoder;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;

import javax.annotation.PostConstruct;
import javax.annotation.PreDestroy;

@Component
@Slf4j
public class WebSocketHandler extends TextWebSocketHandler {
    private final Map<String, WebSocketSession> sessions = new ConcurrentHashMap<>();
    private final Map<String, Long> lastHeartbeatTime = new ConcurrentHashMap<>();
    private final ObjectMapper objectMapper = new ObjectMapper();
    
    // 心跳超时时间（毫秒）
    private static final long HEARTBEAT_TIMEOUT = 30000;
    // 心跳检查间隔（毫秒）
    private static final long CHECK_INTERVAL = 10000;

    @Autowired
    private MessageManager messageManager;

    // 心跳检查调度器
    private final ScheduledExecutorService heartbeatScheduler = Executors.newSingleThreadScheduledExecutor();

    @PostConstruct
    public void init() {
        // 启动心跳检查任务
        heartbeatScheduler.scheduleAtFixedRate(
            this::checkHeartbeats,
            CHECK_INTERVAL,
            CHECK_INTERVAL,
            TimeUnit.MILLISECONDS
        );
    }

    @Override
    public void afterConnectionEstablished(WebSocketSession session) {
        String username = getUsername(session);
        sessions.put(username, session);
        lastHeartbeatTime.put(username, System.currentTimeMillis());
        
        // 广播用户上线消息
        broadcastUserStatus(username, true);
        
        // 发送当前在线用户列表
        sendUserList();
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
        String username = getUsername(session);
        sessions.remove(username);
        lastHeartbeatTime.remove(username);
        
        // 广播用户下线消息
        broadcastUserStatus(username, false);
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage textMessage) {
        try {
            Message message = objectMapper.readValue(textMessage.getPayload(), Message.class);
            
            // 处理心跳消息
            if (message.getType() == Message.Type.HEARTBEAT) {
                handleHeartbeat(session, message);
                return;
            }
            
            // 其他消息交给 MessageManager 处理
            messageManager.handleIncomingMessage(message);
        } catch (Exception e) {
            log.error("Error handling message", e);
        }
    }

    private void handleHeartbeat(WebSocketSession session, Message message) {
        String username = getUsername(session);
        lastHeartbeatTime.put(username, System.currentTimeMillis());
        
        // 发送心跳响应
        try {
            Message response = new Message();
            response.setType(Message.Type.HEARTBEAT);
            response.setTimestamp(System.currentTimeMillis());
            session.sendMessage(new TextMessage(objectMapper.writeValueAsString(response)));
        } catch (Exception e) {
            log.error("Error sending heartbeat response", e);
        }
    }

    private void checkHeartbeats() {
        long now = System.currentTimeMillis();
        List<String> timeoutUsers = new ArrayList<>();

        // 检查所有会话的心跳时间
        lastHeartbeatTime.forEach((username, lastTime) -> {
            if (now - lastTime > HEARTBEAT_TIMEOUT) {
                timeoutUsers.add(username);
            }
        });

        // 关闭超时的会话
        timeoutUsers.forEach(username -> {
            WebSocketSession session = sessions.get(username);
            if (session != null) {
                try {
                    session.close();
                } catch (IOException e) {
                    log.error("Error closing timed out session for user: {}", username, e);
                }
                sessions.remove(username);
                lastHeartbeatTime.remove(username);
                broadcastUserStatus(username, false);
            }
        });
    }

    @PreDestroy
    public void destroy() {
        heartbeatScheduler.shutdown();
        try {
            if (!heartbeatScheduler.awaitTermination(5, TimeUnit.SECONDS)) {
                heartbeatScheduler.shutdownNow();
            }
        } catch (InterruptedException e) {
            heartbeatScheduler.shutdownNow();
            Thread.currentThread().interrupt();
        }
    }

    public void sendMessage(Message message) throws IOException {
        WebSocketSession recipientSession = sessions.get(message.getTo());
        if (recipientSession != null && recipientSession.isOpen()) {
            String messageJson = objectMapper.writeValueAsString(message);
            synchronized (recipientSession) {
                recipientSession.sendMessage(new TextMessage(messageJson));
            }
        } else {
            throw new IOException("Recipient not connected: " + message.getTo());
        }
    }

    private void broadcastUserStatus(String username, boolean isOnline) {
        Message statusMessage = new Message();
        statusMessage.setType(isOnline ? Message.Type.LOGIN : Message.Type.LOGOUT);
        statusMessage.setFrom(username);
        statusMessage.setTimestamp(System.currentTimeMillis());
        
        broadcast(statusMessage);
    }

    private void sendUserList() {
        Message userListMessage = new Message();
        userListMessage.setType(Message.Type.USER_LIST);
        userListMessage.setUsers(new ArrayList<>(sessions.keySet()));
        userListMessage.setTimestamp(System.currentTimeMillis());
        
        broadcast(userListMessage);
    }

    private void broadcast(Message message) {
        try {
            String messageJson = objectMapper.writeValueAsString(message);
            TextMessage textMessage = new TextMessage(messageJson);
            
            for (WebSocketSession session : sessions.values()) {
                if (session.isOpen()) {
                    synchronized (session) {
                        session.sendMessage(textMessage);
                    }
                }
            }
        } catch (IOException e) {
            log.error("Error broadcasting message", e);
        }
    }

    private String getUsername(WebSocketSession session) {
        try {
            String query = session.getUri().getQuery();
            if (query != null && query.startsWith("username=")) {
                String username = query.substring("username=".length());
                return URLDecoder.decode(username, StandardCharsets.UTF_8.name());
            }
            throw new IllegalArgumentException("Username parameter not found in WebSocket URL");
        } catch (Exception e) {
            throw new RuntimeException("Error getting username from session", e);
        }
    }
} 